import rospy
from sensor_msgs.msg import Image
from geometry_msgs.msg import PoseStamped
import cv2
from cv_bridge import CvBridge, CvBridgeError
import numpy as np

# Global variables
targets = []

class Target:
    def __init__(self):
        self.color = 0      # 1: red, 2: green, 3: blue, 4: yellow
        self.shape = 0      # 1: square, 2: circle, 3: triangle
        self.offset_x = 0.0 # up: 0~100 down: -100~0
        self.offset_y = 0.0 # left: 0~100 right: -100~0

def detect_color(frame):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    lower_red1, upper_red1 = np.array([0, 21, 22]), np.array([11, 255, 255])
    lower_red2, upper_red2 = np.array([160, 43, 46]), np.array([180, 255, 255])
    lower_yellow, upper_yellow = np.array([26, 43, 46]), np.array([34, 255, 255])
    lower_green, upper_green = np.array([35, 43, 46]), np.array([77, 255, 255])
    lower_blue, upper_blue = np.array([100, 43, 46]), np.array([124, 255, 255])

    mask_red1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask_red2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask_red = cv2.bitwise_or(mask_red1, mask_red2)
    mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
    mask_green = cv2.inRange(hsv, lower_green, upper_green)
    mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)

    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    total_non_zero = cv2.countNonZero(gray_frame)
    red_non_zero = cv2.countNonZero(mask_red)
    yellow_non_zero = cv2.countNonZero(mask_yellow)
    green_non_zero = cv2.countNonZero(mask_green)
    blue_non_zero = cv2.countNonZero(mask_blue)

    red_ratio = float(red_non_zero) / total_non_zero
    yellow_ratio = float(yellow_non_zero) / total_non_zero
    green_ratio = float(green_non_zero) / total_non_zero
    blue_ratio = float(blue_non_zero) / total_non_zero

    if red_ratio > 0.8:
        return 1
    if yellow_ratio > 0.8:
        return 4
    if green_ratio > 0.8:
        return 2
    if blue_ratio > 0.8:
        return 3

    return 0

def process(frame):
    global targets
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    thresh = cv2.Canny(gray, 69, 200)
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    cv2.imshow("Thresh", thresh)

    for contour in contours:
        area = cv2.contourArea(contour)
        if area < 3000:
            continue

        min_rect = cv2.minAreaRect(contour)
        bbox_area = min_rect[1][0] * min_rect[1][1]
        ratio = area / bbox_area

        M = cv2.moments(contour)
        if M["m00"] != 0:
            cX = int(M["m10"] / M["m00"])
            cY = int(M["m01"] / M["m00"])
        else:
            cX = frame.shape[1] // 2
            cY = frame.shape[0] // 2

        target = Target()
        if ratio > 0.82:
            target.shape = 1  # square
        elif 0.70 < ratio <= 0.82:
            target.shape = 2  # circle
        elif 0.45 < ratio <= 0.55:
            target.shape = 3  # triangle
        else:
            target.shape = 0

        if target.shape:
            mask_roi = np.zeros(frame.shape[:2], dtype=np.uint8)
            cv2.drawContours(mask_roi, [contour], -1, 255, -1)
            framez = cv2.bitwise_and(frame, frame, mask=mask_roi)
            target.color = detect_color(framez)

        target.offset_x = (frame.shape[0] / 2.0 - cY) / frame.shape[0]
        target.offset_y = (frame.shape[1] / 2.0 - cX) / frame.shape[1]

        box = cv2.boxPoints(min_rect)
        box = np.int0(box)
        cv2.drawContours(frame, [box], 0, (0, 255, 0), 2)
        cv2.circle(frame, (cX, cY), 5, (255, 0, 0), -1)

        targets.append(target)

def publish(pub):
    global targets
    for target in targets:
        if not target.shape or not target.color:
            continue

        pose_msg = PoseStamped()
        pose_msg.header.stamp = rospy.Time.now()
        pose_msg.header.frame_id = "camera_coordinate"
        pose_msg.pose.position.x = target.color
        pose_msg.pose.position.y = target.shape
        pose_msg.pose.position.z = target.offset_x
        pose_msg.pose.orientation.w = target.offset_y

        pub.publish(pose_msg)

        rospy.loginfo("Published: [Color:%d, Shape:%d, Offset_X:%f, Offset_Y:%f]",
                      target.color, target.shape, target.offset_x, target.offset_y)
    targets.clear()

def image_callback(msg):
    bridge = CvBridge()
    try:
        frame = bridge.imgmsg_to_cv2(msg, "bgr8")
        if frame is None:
            rospy.logerr("Received empty frame")
            return

        process(frame)
        cv2.imshow("Frame", frame)
        cv2.waitKey(1)
    except CvBridgeError as e:
        rospy.logerr("cv_bridge exception: %s", e)

def main():
    rospy.init_node('vision_node', anonymous=True)
    vision_pub = rospy.Publisher('/ele_vision', PoseStamped, queue_size=10)
    rospy.Subscriber('/usb_cam/image_raw', Image, image_callback)

    rospy.loginfo("Started vision node")

    rate = rospy.Rate(30)
    while not rospy.is_shutdown():
        publish(vision_pub)
        rate.sleep()

if __name__ == '__main__':
    main()